from open_loop_mpc import SolveScheduler
from actual_field import SimulateTheActualPlant
from lstm_model import getRNNmodel_numpy, sigmoidFunction_original
import numpy as np
from matplotlib import pyplot as plt
from parameters_mpc import Sandyclayloam
from richards_equation_1D import volMoistureAllNodes
from van_Genuchten_relations import pressureHead
import time
soilPars = Sandyclayloam()

start_time = time.time()

#Load some relevant data
data_min = np.loadtxt("./weights_casestudy/data_min_scl_c.txt")
data_max = np.loadtxt("./weights_casestudy/data_max_scl_c.txt")
initialState = np.loadtxt("./uncontrolled_inputs/init_state_cl_scl.txt")
initialIrrigation = np.loadtxt("./uncontrolled_inputs/irrig_amount_cl_scl.txt")
initialIrrigation = initialIrrigation*1e07
initialET = np.loadtxt('./uncontrolled_inputs/refEvap_cl_scl.txt')
initialKc = np.loadtxt('./uncontrolled_inputs/crop_coeff_cl_scl.txt')
referenceEvap_all = np.loadtxt("./uncontrolled_inputs/Predicted_refEt.txt")
cropCoefficient_all = np.loadtxt("./uncontrolled_inputs/cropCoeff_all.txt")
allHeadValues = np.loadtxt("./uncontrolled_inputs/All_states_cl_scl.txt")

#Some relevant parameters for the simulation
sequenceLength = 5
horizonLength_orig = 20
horizonLength_weather = int(sequenceLength + horizonLength_orig)
simulationPeriod = 2

#This function computes the root zone pressure head
def compute_root_zone_moisture(Xk):
    root_zone_head = 0.10*(0.25*(Xk[0] + Xk[1] + Xk[2] + Xk[3] + Xk[4])) + 0.20*(0.25*(Xk[4] + Xk[5] + Xk[6] + Xk[7] + Xk[8])) +\
        0.30*(0.25*(Xk[8] + Xk[9] + Xk[10] + Xk[11] + Xk[12])) + 0.40*(0.25*(Xk[12] + Xk[13] + Xk[14] + Xk[15] + Xk[16])) 
    return root_zone_head

#This function converts the root zone pressure head to soil moisture
def obtain_soil_moisture(head):
    vol_moist = volMoistureAllNodes(head, soilPars)
    return vol_moist

def MpcScheduler():
    statesList = []
    irrigationList = []

    for j in range(sequenceLength):
        statesList+=[initialState[j]]
        pass

    for j in range(sequenceLength-1):
        irrigationList+=[initialIrrigation[j]]
        pass

    for j in range(sequenceLength-1):
        referenceEvap_all[j] = initialET[j]
        cropCoefficient_all[j] = initialKc[j]
        pass

    prescribedIrrigation_closedLoop = []
    stateTrajectory_closedLoop = []
    stateTrajectory_openLoop = []
    irrigationDecisions_closedLoop = []
    stateTrajectory_closedLoop.append(initialState[-1])
    stateTrajectory_openLoop.append(initialState[-1])
    allStates=np.zeros((simulationPeriod+1, 17))
    allStates[0,:] = allHeadValues 
    
    #Define the bounds on the state and the slack variables
    xlb = -np.inf*np.ones(1)
    xub = np.inf*np.ones(1)
    slub=np.inf*np.ones(1)
    sllb=np.zeros(1)
    yub = 1
    ylb = -1
    guessY = 0.1

    #Penalty/Weights on the slack variables
    QUpper = 9000
    QLower = 10000
    R = 50
    slope =60
    
    #Initial guesses for the optimization problem 
    xguess_unscaled = 0.22
    xguess = (xguess_unscaled - data_min[0])/(data_max[0] - data_min[0])
    uguess = -1

    #Define the upper and the lower zone
    UpperZone = 0.28
    LowerZone = 0.16

    #Scale the zone
    LZ_scaled = (LowerZone-data_min[0])/(data_max[0] - data_min[0])
    UZ_scaled = (UpperZone-data_min[0])/(data_max[0] - data_min[0])

    #Define the bounds on the input
    UL = -3.9
    UB = -0.04
    guessesX = xguess*np.ones(horizonLength_orig+1)
    guessesU = uguess*np.ones(horizonLength_orig)
    guessesY = guessY*np.ones(horizonLength_orig)
        
    for k in range(simulationPeriod):
        print("\n\n")
        print("-------------------Open-loop simulation for day",k+1)
        start_time_internal = time.time()
    
        x0_init = np.array(statesList[k:k+sequenceLength])
        u_init = np.array(irrigationList[k:k+sequenceLength-1])
        kc_init = cropCoefficient_all[k:k+horizonLength_weather]
        et_init = referenceEvap_all[k:k+horizonLength_weather]

        x0_scaled = (x0_init - data_min[0])/(data_max[0]-data_min[0])
        kc_scaled = (kc_init - data_min[2])/(data_max[2]-data_min[2])
        et_scaled = (et_init - data_min[3])/(data_max[3]-data_min[3])

        irrigationAmounts, irrigationDecisions, states = SolveScheduler(x0_scaled, u_init, kc_scaled, et_scaled,xlb, 
                                      xub, xguess, sllb, slub, QLower, QUpper, LZ_scaled, UZ_scaled, R, uguess, UB, UL, guessY, yub,ylb, slope, guessesX, guessesU, guessesY)
    
    
        modeled_binary = sigmoidFunction_original(slope, np.array(irrigationDecisions))
        print(irrigationAmounts[0])
        
        irrigationList+=[irrigationAmounts[0]]
        prescribedIrrigation_closedLoop+=[irrigationAmounts[0]]        
        stateTrajectory_openLoop+=[states[0]]
        irrigationDecisions_closedLoop+=[modeled_binary[0]]
        
        #Simulate the actual field
        x_currentSE= allStates[k,:]
        u_currentSE = irrigationAmounts[0] * 1e-07 
        kc_currentSE = kc_init[4]
        et_currentSE = et_init[4]
    
        currentStatesSE = SimulateTheActualPlant(x_currentSE, u_currentSE, et_currentSE, kc_currentSE)
        rootzone_head = compute_root_zone_moisture(currentStatesSE)
        rootzone_moist = obtain_soil_moisture(rootzone_head)
        statesList+=[rootzone_moist]
        allStates[k+1,:] = currentStatesSE
    
        #LSTM Model Simulation
        x_current=np.array(statesList[k:k+sequenceLength])
        x_current_scaled = (x_current - data_min[0])/(data_max[0]-data_min[0])
        u_current = np.array(irrigationList[k:k+sequenceLength])
        u_current = u_current * 1e-07
        u_current_scaled = (u_current - data_min[1])/(data_max[1]-data_min[1])
    
        kc_current = kc_scaled[k:k+sequenceLength]
        et_current = et_scaled[k:k+sequenceLength]
        currentState = getRNNmodel_numpy(x_current_scaled, u_current_scaled, kc_current, et_current)
        currentState_unscaled = (currentState * (data_max[0] - data_min[0])) + data_min[0]
        stateTrajectory_closedLoop+=[currentState_unscaled]
    
        elapsed_time_internal = time.time() - start_time_internal
        print(f"\n One open-loop OCP takes {elapsed_time_internal} seconds to solve!\n\n")
    
        #Use the optimal values for the current step as the guesses for the optimization problem in the next time step
        guessesX = (states - data_min[0])/(data_max[0]-data_min[0])
        guessesU = irrigationAmounts
        guessesY = irrigationDecisions
        pass

    elapsed_time = time.time() - start_time 
    print(f"\nThe 20 day simulation takes {elapsed_time} to complete!!\n\n")
    
    #Visualize the results
    
    u_prescribed = np.array(prescribedIrrigation_closedLoop)
    #Obtain the array for the time
    timeStamps = np.arange(1, simulationPeriod+2)
    timeStamps_input = np.arange(1, simulationPeriod+1)

    # #Plot the root zone volumetric moisture content
    # plt.figure(1, figsize=(10,8))
    # plt.plot(timeStamps, np.array(stateTrajectory_closedLoop), color="red", linestyle="dashdot")
    # plt.plot(timeStamps, LowerZone*np.ones(simulationPeriod+1), color = "blue", linestyle = "dashdot")
    # plt.plot(timeStamps, UpperZone*np.ones(simulationPeriod+1), color = "blue", linestyle ="dashdot")
    # plt.ylabel('Volumetric moisture (m3/m3)')
    # plt.xlabel("Time (Days)")
    # plt.xlim(xmin=1)
    # plt.legend(['state Trajectory', 'zone'], loc='upper right') 
    # plt.title('State Trajectory -Closed-Loop')
    # plt.tight_layout()
    # plt.savefig('./test_casestudy_results/state_trajectory_vol.pdf')
    # plt.show()
    # plt.close()

    # plt.figure(2, figsize=(10,8))
    # plt.plot(timeStamps_input, -u_prescribed,marker='o')
    # plt.xlabel("Time (Days)")
    # plt.ylabel("Irrigation amount (mm/day) - Closed-loop")
    # plt.tight_layout()
    # plt.savefig('./test_casestudy_results/control_trajectory_vol.pdf')
    # plt.show()
    # plt.close()


    #Convert the moisture content to pressure head and plot the corresponding results
    LowerZone_head = pressureHead(LowerZone, soilPars)
    UpperZone_head = pressureHead(UpperZone, soilPars)

    stateTrajectory_closedLoop_head = []
    for i in range(len(stateTrajectory_closedLoop)):
        stateTrajectory_closedLoop_head.append(pressureHead(stateTrajectory_closedLoop[i], soilPars))
        pass

    plt.figure(3, figsize=(10,8))
    plt.plot(timeStamps, np.array(stateTrajectory_closedLoop_head), color="red", linestyle="dashdot")
    plt.plot(timeStamps, LowerZone_head*np.ones(simulationPeriod+1), color = "blue", linestyle = "dashdot")
    plt.plot(timeStamps, UpperZone_head*np.ones(simulationPeriod+1), color = "blue", linestyle ="dashdot")
    plt.ylabel('h (m)')
    plt.xlabel("Time (Days)")
    plt.xlim(xmin=1)
    plt.legend(['state Trajectory', 'zone'], loc='upper right') 
    plt.title('State Trajectory -Closed-Loop')
    plt.tight_layout()
    plt.savefig('./test_casestudy_results/state_trajectory_head.pdf')
    
    plt.figure(4, figsize=(10,8))
    plt.plot(timeStamps_input, -u_prescribed,marker='o')
    plt.xlabel("Time (Days)")
    plt.ylabel("Irrigation amount (mm/day) - Closed-loop")
    plt.tight_layout()
    plt.savefig('./test_casestudy_results/control_trajectory_head.pdf')
    plt.show()
    plt.close()
    
    return 

if __name__ == '__main__':
    MpcScheduler()
    

  
